{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbed92dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import pandas as pd\n",
    "import os\n",
    "import numpy as np\n",
    "from glob import glob\n",
    "from onekey_algo import get_param_in_cwd\n",
    "from onekey_algo.custom.components.metrics import calc_dice, calc_iou\n",
    "\n",
    "details = []\n",
    "os.makedirs('data', exist_ok=True)\n",
    "os.makedirs('results', exist_ok=True)\n",
    "for model_name in os.listdir(get_param_in_cwd('save_dir')):\n",
    "    samples = glob(os.path.join(get_param_in_cwd('data_root'), '*', 'masks', r'*[jpg|png|bmp]'))\n",
    "    print(f\"一共获取到{len(samples)}\")\n",
    "    for sample in samples:\n",
    "        cohort = os.path.basename(os.path.dirname(os.path.dirname(sample)))\n",
    "        simg = Image.open(sample)\n",
    "        infer = os.path.join(os.path.join(get_param_in_cwd('save_dir', '.'), get_param_in_cwd('model_name'), 'test_results', 'masks'),\n",
    "                             os.path.basename(sample))\n",
    "        iimg = Image.open(infer)\n",
    "        dice = calc_dice(np.array(iimg), np.array(simg))\n",
    "        iou = calc_iou(np.array(iimg), np.array(simg))\n",
    "        details.append([os.path.basename(sample), dice, iou, model_name, cohort])\n",
    "details = pd.DataFrame(details, columns=['sample', 'Dice', 'mIoU', 'Model', 'Cohort'])\n",
    "details.to_csv('results/metrics_details.csv', index=False)\n",
    "details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ba4ddc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = details.groupby(['Model', 'Cohort']).aggregate('mean').reset_index()\n",
    "metrics.to_csv('results/metrics.csv', index=False)\n",
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00c41b64",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
